# 添加父级目录到路径
import os
import sys
import numpy as np
import faiss
current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(current_dir)
print(f"current_dir: {current_dir}")
sys.path.insert(0, parent_dir)
from my_common import load_auto_model, load_auto_tokenizer, load_flag_model


# Step 1: Dataset
corpus = [
    "迈克尔杰克逊是一位传奇的流行偶像，以其破纪录的音乐和舞蹈创新而闻名。",
    "Fei-Fei Li 是斯坦福大学的教授，他通过 ImageNet 项目彻底改变了计算机视觉。",
    "布拉德·皮特 （Brad Pitt） 是一位多才多艺的演员和制片人，以其在《搏击俱乐部》和《好莱坞往事》等电影中的角色而闻名。",
    "Geoffrey Hinton 作为 AI 的奠基人物，因其在深度学习方面的贡献而获得图灵奖。",
    "Eminem 是一位著名的说唱歌手，也是有史以来最畅销的音乐艺术家之一。",
    "泰勒·斯威夫特 （Taylor Swift） 是一位获得格莱美奖的创作歌手，以其叙事驱动的音乐而闻名。",
    "山姆·阿尔特曼 （Sam Altman） 作为 OpenAI 的首席执行官，以 GPT 系列的惊人作品，追求安全有益的 AI。",
    "摩根·弗里曼 （Morgan Freeman） 是一位备受赞誉的演员，以其独特的声音和多样化的角色而闻名。",
    "Andrew Ng 通过 Coursera 和斯坦福大学的公开课程在全球范围内传播 AI 知识。",
    "小罗伯特·唐尼 （Robert Downey Jr.） 是一位标志性的演员，以在漫威电影宇宙中扮演钢铁侠而闻名。"
]
queries = [
    "小罗伯特·唐尼是谁？",
    "神经网络专家",
    "一位著名的女歌手",
]

if __name__ == '__main__':
    flag_model = load_flag_model(query_instruction_for_retrieval='Represent this sentence for searching relevant passages:')
    # Step 2: Text Embedding
    # get the embedding of the corpus
    corpus_embeddings = flag_model.encode(corpus)

    print("shape of the corpus embeddings:", corpus_embeddings.shape)
    print("data type of the embeddings: ", corpus_embeddings.dtype)

    # Faiss only accepts float32 inputs
    corpus_embeddings = corpus_embeddings.astype(np.float32)

    # Step 3: Indexing
    # get the length of our embedding vectors, vectors by bge-base-en-v1.5 have length 768；bge-large-zh-v1.5 have length 1024
    dim = corpus_embeddings.shape[-1]

    # 创建 Faiss 索引并将语料库嵌入存储到向量空间中
    index = faiss.index_factory(dim, 'Flat', faiss.METRIC_INNER_PRODUCT)

    # 如果我们使用 "Flat" 量化器并将 METRIC_INNER_PRODUCT 作为度量，则无需训练。其他一些使用量化的索引可能需要训练
    # check if the index is trained
    print(index.is_trained)  
    # index.train(corpus_embeddings)

    # add all the vectors to the index
    index.add(corpus_embeddings)

    print(f"total number of vectors: {index.ntotal}")

    # Step 3.5 (Optional): Saving Faiss index
    # change the path to where you want to save the index
    index_file_path = f"{current_dir}/faiss_index.bin"
    print(f"index type: {type(index)}")
    # 注意： gpu索引 IndexReplicas 由多个独立子索引组成，不能直接使用write_index保存。所以不能先转到gpu再保存
    faiss.write_index(index, index_file_path)

    # If you already have stored index in your local directory, you can load it by:
    index = faiss.read_index(index_file_path)

     # 如果您安装了 faiss-gpu，可以把索引加载到gup上。
    co = faiss.GpuMultipleClonerOptions()
    index = faiss.index_cpu_to_all_gpus(index, co)

    # Step 4: Find answers to the query
    # get the embeddings of all the queries:
    query_embeddings = flag_model.encode_queries(queries)
    # Then, use the Faiss index to do a knn search in the vector space:
    dists, ids = index.search(query_embeddings, k=3)
    print("search result dists: ")
    print(dists)
    print("search result ids: ")
    print(ids)

    # Now let's see the result:
    for i, q in enumerate(queries):
        print(f"query:\t{q}\nanswer:\t{corpus[ids[i][0]]}\n")

